import warnings
warnings.filterwarnings('ignore')
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei'] # 绘图显示中文
from model import Arima

DATA_CSV_PATH = 'data.csv'
CO2_DATA_CSV_PATH = 'co2_data.csv'

# 读取数据
df = pd.read_excel("./dataDay.xlsx")
df = df[df["sector"]=="Total"]
df = df[["date","co2"]]
df["date"] = df["date"].apply(lambda x:(str(x)[:7]+"-01"))
df['date'] = pd.to_datetime(df['date'], format='%Y-%m-%d')
df = df.groupby("date").sum()

# 数据划分
train_data = df[df.index < '2023-01-01']
test_data = df[df.index >= '2023-01-01']

# 训练模型
arima = Arima()
results = arima.train(train_data=train_data)

# 保存模型
from joblib import dump
dump(results, 'arima_model.joblib')

# 预测并计算相对误差
pred = results.get_forecast('2030-12-01')
actual_values = test_data['co2']
predicted_dates = pd.date_range(start='2023-01-01', end='2030-12-01', freq='MS')
predicted_values = pd.Series(pred.predicted_mean, index=predicted_dates)

# 计算相对误差
monthly_relative_errors = pd.DataFrame(index=predicted_dates, columns=['Relative Error'])
for date, actual in actual_values.items():
    predicted = predicted_values.loc[date]
    if actual != 0:
        relative_error = np.abs(predicted - actual) / np.abs(actual)
        monthly_relative_errors.loc[date, 'Relative Error'] = relative_error
    else:
        monthly_relative_errors.loc[date, 'Relative Error'] = np.nan

print(monthly_relative_errors)

# 保存误差数据到CSV文件
error_data = {
    'date': predicted_dates,
    'actual_value': actual_values,
    'predicted_value': predicted_values,
    'relative_error': monthly_relative_errors['Relative Error']
}
error_df = pd.DataFrame(error_data)
error_df.to_csv(DATA_CSV_PATH, index=False)

# 保存预测数据到CO2数据文件
co2_data = {
    'date': predicted_dates,
    'predicted_value': predicted_values
}
co2_df = pd.DataFrame(co2_data)
co2_df.to_csv(CO2_DATA_CSV_PATH, index=False)

# 绘图
ax = df.plot()
pred.predicted_mean['2023-01-01':'2024-09-01'].plot(ax=ax, label='predict value',color="red")
plt.title("China Residential CO2 Index")
plt.ylabel('MtCO2 Per Month')
plt.xlabel('Date')
plt.legend()
plt.savefig('static/result.png')
